import requests
import base64
import datetime
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class TwistChecker:
    def __init__(self, api_url: str = "http://localhost:11434/api/chat",
                 model: str = "qwen2.5vl:3b"):
        """
        初始化 Ollama 客户端

        参数:
            api_url: API 服务地址
            model: 默认使用的模型名称
        """
        self.api_url = api_url
        self.model = model
        self.headers = {"Content-Type": "application/json"}

    def check(self, pic_path: str,
              prompt: str = None,
              model: str = None) -> dict:
        """
        检查图片中的铁丝缠绕情况

        参数:
            pic_path: 图片文件路径
            prompt: 自定义提示词，如果为None则使用默认提示词
            model: 自定义模型，如果为None则使用实例默认模型

        返回:
            dict: 包含检查结果的字典，格式为:
                {
                    'success': bool,
                    'content': str,
                    'error': str,
                    'start_time': str,
                    'end_time': str,
                    'duration': float
                }
        """
        # 默认提示词
        if prompt is None:
            prompt = "描述这张图片中是否有铁丝缠绕？如果没有，直接回复没有。如果有，请描述缠绕数量的多少。"

        # 使用指定模型或默认模型
        use_model = model or self.model

        # 记录开始时间
        start_time = datetime.datetime.now()
        start_time_str = start_time.strftime("%Y-%m-%d %H:%M:%S")

        logger.info("开始检查铁丝缠绕情况，开始时间: %s", start_time_str)

        result = {
            'success': False,
            'content': '',
            'error': '',
            'start_time': start_time_str,
            'end_time': '',
            'duration': 0.0
        }

        try:
            # 构建请求载荷
            payload = {
                "model": use_model,
                "messages": [
                    {
                        "role": "user",
                        "content": prompt,
                        "images": [self._image_to_base64(pic_path)]
                    }
                ],
                "stream": False
            }

            # 发送请求
            logger.info("正在调用API...")
            response = requests.post(self.api_url, json=payload, headers=self.headers)

            # 记录结束时间
            end_time = datetime.datetime.now()
            end_time_str = end_time.strftime("%Y-%m-%d %H:%M:%S")
            duration = (end_time - start_time).total_seconds()

            result.update({
                'end_time': end_time_str,
                'duration': duration
            })

            # 处理响应
            if response.status_code == 200:
                response_data = response.json()
                reply_content = response_data.get("message", {}).get("content", "")

                result.update({
                    'success': True,
                    'content': reply_content
                })

                logger.info("模型回复: %s", reply_content)
                logger.info("检查完成，耗时: %.2f 秒", duration)

            else:
                error_msg = f"API 请求失败，状态码: {response.status_code}, 响应: {response.text}"
                result['error'] = error_msg
                logger.error(error_msg)

        except requests.exceptions.RequestException as e:
            error_msg = f"网络请求异常: {str(e)}"
            result['error'] = error_msg
            logger.error(error_msg)

        except Exception as e:
            error_msg = f"调用异常: {str(e)}"
            result['error'] = error_msg
            logger.error(error_msg)

        return result

    def _image_to_base64(self, image_path: str) -> str:
        """
        将图片文件转换为 Base64 编码字符串

        参数:
            image_path: 图片文件路径

        返回:
            str: Base64 编码的图片数据

        异常:
            FileNotFoundError: 图片文件不存在
            Exception: 图片读取或编码失败
        """
        try:
            with open(image_path, "rb") as image_file:
                encoded = base64.b64encode(image_file.read()).decode("utf-8")
                logger.debug("图片 %s 已成功转换为 Base64 编码", image_path)
                return encoded
        except FileNotFoundError:
            error_msg = f"图片文件不存在: {image_path}"
            logger.error(error_msg)
            raise FileNotFoundError(error_msg)
        except Exception as e:
            error_msg = f"图片读取失败: {str(e)}"
            logger.error(error_msg)
            raise Exception(error_msg)


# 便捷函数
def check(pic_path: str, prompt: str = None) -> dict:
    """
    快速检查图片中的铁丝缠绕情况

    参数:
        pic_path: 图片文件路径
        prompt: 自定义提示词

    返回:
        dict: 检查结果
    """
    return TwistChecker().check(pic_path, prompt)